// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <stdint.h>
#include <string.h>

#include <melon/executors/cpu_thread_pool_executor.h>
#include <pollux/common/base/prefix_sort_config.h>
#include <pollux/common/compression/compression.h>

namespace kumo::pollux::common {
#define POLLUX_SPILL_LIMIT_EXCEEDED(errorMessage)                    \
  _POLLUX_THROW(                                                     \
      ::kumo::pollux::PolluxRuntimeError,                         \
      ::kumo::pollux::error_source::kErrorSourceRuntime.c_str(), \
      ::kumo::pollux::error_code::kSpillLimitExceeded.c_str(),   \
      /* isRetriable */ true,                                       \
      "{}",                                                         \
      errorMessage);

#define POLLUX_GENERIC_SPILL_FAILURE(errorMessage)                   \
  _POLLUX_THROW(                                                     \
      ::kumo::pollux::PolluxRuntimeError,                         \
      ::kumo::pollux::error_source::kErrorSourceRuntime.c_str(), \
      ::kumo::pollux::error_code::kGenericSpillFailure.c_str(),  \
      /* isRetriable */ true,                                       \
      "{}",                                                         \
      errorMessage);

    /// Defining type for a callback function that returns the spill directory path.
    /// Implementations can use it to ensure the path exists before returning.
    using GetSpillDirectoryPathCB = std::function<std::string_view()>;

    /// The callback used to update the aggregated spill bytes of a query. If the
    /// query spill limit is set, the callback throws if the aggregated spilled
    /// bytes exceed the set limit.
    using UpdateAndCheckSpillLimitCB = std::function<void(uint64_t)>;

    /// Specifies the config for spilling.
    struct SpillConfig {
        SpillConfig() = default;

        SpillConfig(
            GetSpillDirectoryPathCB _getSpillDirPathCb,
            UpdateAndCheckSpillLimitCB _updateAndCheckSpillLimitCb,
            std::string _filePath,
            uint64_t _maxFileSize,
            uint64_t _writeBufferSize,
            uint64_t _readBufferSize,
            melon::Executor *_executor,
            int32_t _minSpillableReservationPct,
            int32_t _spillableReservationGrowthPct,
            uint8_t _startPartitionBit,
            uint8_t _numPartitionBits,
            int32_t _maxSpillLevel,
            uint64_t _maxSpillRunRows,
            uint64_t _writerFlushThresholdSize,
            const std::string &_compressionKind,
            std::optional<PrefixSortConfig> _prefixSortConfig = std::nullopt,
            const std::string &_fileCreateConfig = {});

        /// Returns the spilling level with given 'startBitOffset' and
        /// 'numPartitionBits'.
        ///
        /// NOTE: we advance (or right shift) the partition bit offset when goes to
        /// the next level of recursive spilling.
        int32_t spillLevel(uint8_t startBitOffset) const;

        /// Checks if the given 'startBitOffset' has exceeded the max spill limit.
        bool exceedSpillLevelLimit(uint8_t startBitOffset) const;

        /// Returns true if prefix sort is enabled.
        bool prefixSortEnabled() const {
            return prefixSortConfig.has_value();
        }

        /// A callback function that returns the spill directory path. Implementations
        /// can use it to ensure the path exists before returning.
        GetSpillDirectoryPathCB getSpillDirPathCb;

        /// The callback used to update the aggregated spill bytes of a query. If the
        /// query spill limit is set, the callback throws if the aggregated spilled
        /// bytes exceed the set limit.
        UpdateAndCheckSpillLimitCB updateAndCheckSpillLimitCb;

        /// Prefix for spill files.
        std::string fileNamePrefix;

        /// The max spill file size. If it is zero, there is no limit on the spill
        /// file size.
        uint64_t maxFileSize;

        /// Specifies the size to buffer the serialized spill data before write to
        /// storage system for io efficiency.
        uint64_t writeBufferSize;

        /// Specifies the buffer size to read from one spilled file. If the underlying
        /// filesystem supports async read, we do read-ahead with double buffering,
        /// which doubles the buffer used to read from each spill file.
        uint64_t readBufferSize;

        /// Executor for spilling. If nullptr spilling writes on the Driver's thread.
        melon::Executor *executor; // Not owned.

        /// The minimal spillable memory reservation in percentage of the current
        /// memory usage.
        int32_t minSpillableReservationPct;

        /// The spillable memory reservation growth in percentage of the current
        /// memory usage.
        int32_t spillableReservationGrowthPct;

        /// The start partition bit offset of the top (the first level) partitions.
        uint8_t startPartitionBit;

        /// Used to calculate the spill hash partition number for hash join and
        /// RowNumber with 'startPartitionBit'.
        uint8_t numPartitionBits;

        /// The max allowed spilling level with zero being the initial spilling
        /// level. This only applies for hash build spilling which needs recursive
        /// spilling when the build table is too big. If it is set to -1, then there
        /// is no limit and then some extreme large query might run out of spilling
        /// partition bits at the end.
        int32_t maxSpillLevel;

        /// The max row numbers to fill and spill for each spill run. This is used to
        /// cap the memory used for spilling. If it is zero, then there is no limit
        /// and spilling might run out of memory.
        uint64_t maxSpillRunRows;

        /// Minimum memory footprint size required to reclaim memory from a file
        /// writer by flushing its buffered data to disk.
        uint64_t writerFlushThresholdSize;

        /// CompressionKind when spilling, CompressionKind_NONE means no compression.
        common::CompressionKind compressionKind;

        /// Prefix sort config when spilling, enable prefix sort when this config is
        /// set, otherwise, fallback to timsort.
        std::optional<PrefixSortConfig> prefixSortConfig;

        /// Custom options passed to pollux::FileSystem to create spill WriteFile.
        std::string fileCreateConfig;
    };
} // namespace kumo::pollux::common
